import re
from typing import Dict, List

from pydantic import BaseModel, Field


def parse_template_from_markdown_text(markdown_text):
    markdown_text = markdown_text.strip()
    lines = markdown_text.split("\n")

    sections = []
    current_section = {"title": None, "subsections": [], "content": "", "raw_str": ""}
    current_subsection = None

    start = 0
    for i, line in enumerate(lines):
        if re.match(r"^# ", line):  # Match level 1 heading
            # Save the previous section and start a new one
            if current_section["title"] or current_section["content"]:
                if current_subsection:
                    current_section["subsections"].append(current_subsection)
                    current_subsection = None
                current_section["raw_str"] = "\n".join(lines[start + 1:i]).strip()
                start = i
                sections.append(current_section)
                current_section = {"title": None, "subsections": [], "content": ""}
            current_section["title"] = line[2:]
        elif re.match(r"^## ", line):  # Match level 2 heading
            # Save the previous subsection and start a new one
            if current_subsection:
                current_section["subsections"].append(current_subsection)
            current_subsection = {"title": line[3:], "content": ""}
        else:
            if current_subsection:
                current_subsection["content"] += line + "\n"
            else:
                current_section["content"] += line + "\n"

    # Append the last section and subsection if they exist
    if current_subsection:
        current_section["subsections"].append(current_subsection)
    if current_section["title"] or current_section["content"]:
        current_section["raw_str"] = "\n".join(lines[start + 1:]).strip()
        sections.append(current_section)

    return {section["title"]: section for section in sections}


class PromptTemplate(BaseModel):
    class Section(BaseModel):
        class Subsection(BaseModel):
            title: str
            content: str

        title: str
        content: str
        raw_str: str
        subsections: List[Subsection] = []

    system_prompt_template: Section = Field(None, alias="System Prompt")
    polish_user_prompt_template: Section = Field(None, alias="Polish User Prompt")
    reply_user_prompt_template: Section = Field(None, alias="Reply User Prompt")
    refine_user_prompt_template: Section = Field(None, alias="Refine User Prompt")

    @classmethod
    def from_markdown_text(cls, markdown_text):
        sections = parse_template_from_markdown_text(markdown_text)
        return cls.model_validate(sections)


class Prompt(BaseModel):
    template: PromptTemplate
    lines: Dict[str, List[str]] = {}
    placeholders: Dict[str, str] = {}

    @property
    def system_prompt(self):
        return self.template.system_prompt_template.format(**self.placeholders)

    @property
    def polish_user_prompt(self):
        return self.template.polish_user_prompt_template.format(**self.placeholders)

    @property
    def reply_user_prompt(self):
        return self.template.reply_user_prompt_template.format(**self.placeholders)

    @property
    def refine_user_prompt(self):
        return self.template.refine_user_prompt_template.format(**self.placeholders)
